import numpy as np
from lte import constant as c
from lte.phy import Phy
from lte.sequence import Cell

class refsignal(Cell):
    offset = [10,558,970,1518]
    symbol_idx = [0,4,7,11]
    def __init__(self,id,nof_prb=6,cp="normal"):
        Cell.__init__(self,id,nof_prb,cp)
        self.symbol_sz = c.LTE_SYMBOL_SZ[self.nof_prb]
        self.sf_len = 1920
        
        self.ref = np.zeros((self.nof_ports,c.LTE_NOF_SF_X_FRAME,self.sf_len),dtype=complex)
        self.frame_ref_freq = np.zeros((self.nof_ports,c.LTE_NOF_SF_X_FRAME*self.sf_len),dtype=complex)
        
        for p in range(self.nof_ports):
            self.ref[p,:,:] = self.cs_tx(p)
        for p in range(self.nof_ports):
            self.frame_ref_freq[p,:] = np.fft.fft(self.frame_refsignal(p))

    def xcorr(self,port_id,buf,slot,peak):
        ch = np.zeros(128,dtype=complex)
        for off in self.offset:
            o = off+peak
            r = self.ref[port_id][slot,off:off+self.symbol_sz]
            for i in range(-10,64):
                d = buf[o+i:o+i+self.symbol_sz]
                ch[i+64] += np.dot(d,np.conj(r))
        return ch

    def xcorr_slot(self,port_id,buf,slot,peak):
        ch = np.zeros(128,dtype=complex)
        r = self.ref[port_id][slot,:]
        for i in range(-10,64):
            d = buf[peak+i:peak+i+c.LTE_SF_LEN(self.symbol_sz)]
            ch[i+64] += np.dot(d,np.conj(r))
        return ch

    def pos0(self,port_id,buf,peak):
        res = np.zeros((c.LTE_NOF_SF_X_FRAME,),dtype=complex)
        sf_idx_init = (2 * c.LTE_NOF_SF_X_FRAME - peak // self.sf_len) % c.LTE_NOF_SF_X_FRAME
        n      = peak % self.sf_len
        for _slot in range(c.LTE_NOF_SF_X_FRAME):
            slot = (sf_idx_init+_slot)%c.LTE_NOF_SF_X_FRAME
            for off in self.offset:
                d = buf[off+n:off+n+self.symbol_sz]
                r = self.ref[port_id][slot,off:off+self.symbol_sz]
                res[_slot] += np.dot(d,np.conj(r))
            n += self.sf_len
        return res

    def frame(self,port_id,buf,peak):
        ch = np.zeros(128,dtype=complex)
        sf_idx_init = (2 * 10 - peak // self.sf_len) % 10
        n      = peak % self.sf_len
        for _slot in range(10):
            slot = (sf_idx_init+_slot)%10
            ch += self.xcorr_slot(port_id,buf,slot,n)
            n += self.sf_len
        return ch

    def symbol_of_frame(self,port_id,buf,peak):
        ch = np.zeros(19200,dtype=complex)
        sf_idx_init = (2 * 10 - peak // self.sf_len) % 10
        n      = peak % self.sf_len
        for _slot in range(10):
            slot = (sf_idx_init+_slot)%10
            sym =self.symbol_of_slot(port_id,buf,slot,n)
            for i in range(4):
                ch[_slot*1920+self.offset[i]] = sym[i] 
            n += self.sf_len
        return ch
    
    def symbol_of_slot(self,port_id,buf,slot,peak):
        ch = np.zeros(4,dtype=complex)
        r = self.ref[port_id][slot,:]
        for i in range(4):
            off = self.offset[i]
            d = buf[off+peak:off+peak+self.symbol_sz]
            r = self.ref[port_id][slot,off:off+self.symbol_sz]
            ch[i] = np.dot(d,np.conj(r))
        return ch

    def afc(self,buf,port,peak):
        l = len(buf)
        p = (l-peak)//19200
        r = np.zeros(p*10-10,dtype=complex)
        for frame in range(p-1):
            r[frame*10:frame*10+10] = self.pos0(port,buf[frame*19200:frame*19200+19200*2],peak)
        return r

    def frame_refsignal(self,port_id):
        r = np.zeros(19200,dtype=complex)
        for slot in range(10):
            o = slot*1920 
            r[o:o+1920] = self.ref[port_id][slot,:]
        return r

    def find_freq(self,buf,port_id,frange):
        d = buf[:19200]
        fd = np.fft.fft(d)
        flen = len(frange)
        txr = np.zeros((flen,19200))
        for i in range(flen):
            s = frange[i]
            fr = np.zeros(19200,dtype=complex)
            if s !=0:
                fr[s:] = self.frame_ref_freq[port_id][:-s]
                fr[:s] = self.frame_ref_freq[port_id][-s:]
            else:
                fr[:] = self.frame_ref_freq[port_id]
            xr = fd*np.conj(fr)
            txr[i,:] = np.abs(np.fft.ifft(xr))
        peak = txr.argmax()
        df = peak//19200
        peak = peak%19200
        return txr[df,peak],port_id,peak,frange[df]

    def find_peak(self,buf,port_id):
        d = buf[:19200]
        fd = np.fft.fft(d)
        xr = fd*np.conj(self.frame_ref_freq[port_id])
        txr = np.abs(np.fft.ifft(xr))
        peak = txr.argmax()
        return txr[peak],peak,port_id

    def find_all_peak(self,buf):
        r = [self.find_peak(buf,port_id) for port_id in range(self.nof_ports)]
        r.sort(reverse=True)
        return r

    def find_all_freq(self,buf,frange):
        r = [self.find_freq(buf,port_id,frange) for port_id in range(self.nof_ports)]
        r.sort(reverse=True)
        return r

    @staticmethod
    def all_peak(buf):
        r = []
        for cid in range(504):
            for port in range(4):
                ref = refsignal(cid)
                a,peak,_ = ref.find_peak(buf,port)
                r.append((a,cid,port,peak))
        r.sort(reverse=True)
        return r

class LazyReferences:
    def __init__(self):
        self.refs = [None] * 504
    
    def __getitem__(self, idx):
        if self.refs[idx]==None:
            self.refs[idx] = refsignal(idx)
        return self.refs[idx]
